from omegaconf import DictConfig
from torchvision.transforms.v2 import Compose, RandomApply

from ecgcmr.signal.sig_augmentations.RandomCrop import RandomCropSignal
from ecgcmr.signal.sig_augmentations.FTSurrogate import FTSurrogate
from ecgcmr.signal.sig_augmentations.Jitter import Jitter
from ecgcmr.signal.sig_augmentations.Rescaling import Rescaling
from ecgcmr.signal.sig_augmentations.ShuffleChannels import ChannelShuffle
from ecgcmr.signal.sig_augmentations.Flipping import TimeReverse, SignFlip

from ecgcmr.multimodal.multimodal_utils.ToTensor import ToTensor
from ecgcmr.multimodal.multimodal_utils.SqueezeUnsqueeze import UnsqueezeDimension


class ECGAugmentations:
    def __init__(
            self,
            cfg: DictConfig,
            sampling_rate: int = 500,
            input_electrodes: int = 12,
            single_modality: bool = True,
            seed: int = 2024,
            use_peaks_location: bool = False,
            apply_augmentations: bool = True,
        ):
        self.sampling_rate = sampling_rate
        self.input_electrodes = input_electrodes
        self.single_modality = single_modality
        self.seed = seed
        self.apply_augmentations = apply_augmentations
        self.ecg_aug_cfg = cfg
        self.use_peaks_location = use_peaks_location

        self.augmentation_list = []

        self._init_augmentations()

    def _init_augmentations(self):
        self._add_random_crop()
        if self.apply_augmentations:
            self._add_ft_surrogate()
            self._add_jitter()
            self._add_rescaling()
            self._add_channel_shuffle()
            self._add_time_reverse()
            self._add_sign_flip()
        self._add_to_tensor()
        self._add_unsqueeze_dimension()
        
        self.ecg_augmentation_steps = Compose([aug for aug in self.augmentation_list if aug is not None])

    def _add_ft_surrogate(self):
        if self.ecg_aug_cfg.ft_surrogate.enable:
            keys = None if self.single_modality else ["ecg"]
            self.augmentation_list.append(
                FTSurrogate(
                    phase_noise_magnitude=self.ecg_aug_cfg.ft_surrogate.phase_noise_magnitude,
                    channel_indep=self.ecg_aug_cfg.ft_surrogate.channel_indep,
                    seed=self.seed,
                    prob=self.ecg_aug_cfg.ft_surrogate.prob,
                    keys=keys
                )
            )

    def _add_jitter(self):
        if self.ecg_aug_cfg.jitter.enable:
            keys = None if self.single_modality else ["ecg"]
            self.augmentation_list.append(
                RandomApply(
                    [Jitter(
                        sigma=self.ecg_aug_cfg.jitter.sigma,
                        amplitude=self.ecg_aug_cfg.jitter.amplitude,
                        keys=keys
                    )],
                p=self.ecg_aug_cfg.jitter.prob)
            )

    def _add_rescaling(self):
        if self.ecg_aug_cfg.rescaling.enable:
            keys = None if self.single_modality else ["ecg"]
            self.augmentation_list.append(
                RandomApply(
                    [Rescaling(
                        sigma=self.ecg_aug_cfg.rescaling.sigma,
                        keys=keys
                    )],
                    p=self.ecg_aug_cfg.rescaling.prob)
            )

    def _add_random_crop(self):
        if self.ecg_aug_cfg.random_crop.enable:
            self.augmentation_list.append(
                RandomCropSignal(crop_size=self.ecg_aug_cfg.random_crop.ecg_time_steps)
            )

    def _add_channel_shuffle(self):
        if self.ecg_aug_cfg.channel_shuffle.enable:
            keys = None if self.single_modality else ["ecg"]
            self.augmentation_list.append(
                RandomApply(
                    [ChannelShuffle(
                        prob=self.ecg_aug_cfg.channel_shuffle.prob,
                        keys=keys)],
                    p=self.ecg_aug_cfg.channel_shuffle.prob)
            )

    def _add_time_reverse(self):
        if self.ecg_aug_cfg.time_reverse.enable:
            keys = None if self.single_modality else ["ecg", "rpeaks"]
            self.augmentation_list.append(
                RandomApply(
                    [TimeReverse(keys=keys)],
                    p=self.ecg_aug_cfg.time_reverse.prob
                )
            )

    def _add_sign_flip(self):
        if self.ecg_aug_cfg.sign_flip.enable:
            keys = None if self.single_modality else ["ecg"]
            self.augmentation_list.append(
                RandomApply(
                    [SignFlip(keys=keys)],
                    p=self.ecg_aug_cfg.sign_flip.prob
                )
            )

    def _add_to_tensor(self):
        if self.single_modality:
            keys = None
        elif self.use_peaks_location:
            keys = ["ecg", "rpeaks"]
        else:
            keys = ["ecg"]

        self.augmentation_list.append(
            ToTensor(keys=keys)
        )

    def _add_unsqueeze_dimension(self):
        if self.ecg_aug_cfg.unsqueeze.enable:
            keys = None if self.single_modality else ["ecg"]
            self.augmentation_list.append(
                UnsqueezeDimension(
                    position=self.ecg_aug_cfg.unsqueeze.position,
                    keys=keys
                )
            )

    def __call__(self, data: dict) -> dict:
        return self.ecg_augmentation_steps(data)
